from gensim import models
import pprint
from collections import defaultdict
from gensim import corpora
from gensim import similarities
from apps.neo4j.gensim_train.corpus import corpus_raw as documents
import jieba

# ————————————————————————step1————————————————————————
# 原始语句
lst_raw = documents

real_documents = [list(jieba.cut(item_text, cut_all=False)) for item_text in lst_raw]
print(real_documents)

# 语料库
lst = [' '.join(ele) for ele in real_documents]



# ————————————————————————step2————————————————————————
# 数据的预处理
'''

    如果需要删除一些不重要的词汇
'''
stoplist = set(', . ? ! ; : \' " ， 。 ？ ！ ~ ( ) “ ” 、'.split(' '))
#print(stoplist)
texts = [[word for word in document.lower().split() if word not in stoplist]
         for document in lst]

#print(texts)


# ————————————————————————step3————————————————————————
# 词袋方式
# 统计频率
frequency = defaultdict(int)
for text in texts:
    for token in text:
        frequency[token] += 1

# 只保存出现频率大于等于1的
processed_corpus = [[token for token in text if frequency[token] >= 1] for text in texts]
pprint.pprint(processed_corpus)


dictionary = corpora.Dictionary(processed_corpus)
# dictionary.save("my_dict.dict")
print(dictionary)
num_features = len(dictionary)
print(dictionary.token2id)

bow_corpus = [dictionary.doc2bow(text) for text in processed_corpus]
pprint.pprint(bow_corpus)

# ————————————————————————step4————————————————————————
print("————————————————————————————训练模型————————————————————————————")
print("为相似性查询做准备")
tfidf = models.TfidfModel(bow_corpus, dictionary = dictionary)
print(dictionary.__dict__)

print(tfidf.id2word)
print(tfidf[bow_corpus])

index = similarities.SparseMatrixSimilarity(tfidf[bow_corpus], num_features=num_features)



# 保存模型
#tfidf.save("my_model.tfi")
#index.save('my_index.index')
#models.TfidfModel.load("my_model.tif")





# 读取模型
#index = similarities.MatrixSimilarity.load('/tmp/deerwester.index')




# 相似度查询
query_document = "减速机 振动 声音 大".split()
query_bow = dictionary.doc2bow(query_document)

sims = index[tfidf[query_bow]]

#print(list(enumerate(sims)))


for document_number, score in sorted(enumerate(sims), key = lambda x:x[1],reverse = True):
    print(document_number, score)


